import sys
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QLineEdit, QPushButton, QFileDialog, QProgressBar
from PyQt5.QtCore import QThread, pyqtSignal, pyqtSlot
from m3u8handle import download_ts_files
from multiprocessing import Value

class DownloadThread(QThread):
    progress_signal = pyqtSignal(int, int)  # 发射当前进度和总进度
    finished_signal = pyqtSignal()

    def __init__(self, m3u8_url, output_folder, cookies, proxies):
        super().__init__()
        self.m3u8_url = m3u8_url
        self.output_folder = output_folder
        self.cookies = cookies
        self.proxies = proxies
        self.is_running = Value('i', 1)

    def run(self):
        try:
            download_ts_files(
                self.m3u8_url,
                self.output_folder,
                self.progress_callback,
                self.is_running,
                self.cookies,
                self.proxies
            )
        finally:
            self.finished_signal.emit()

    def progress_callback(self, current, total):
        self.progress_signal.emit(current, total)

    @pyqtSlot()
    def request_stop(self):
        self.is_running.value = 0

class App(QWidget):
    def __init__(self):
        super().__init__()
        self.thread = None
        self.initUI()

    def initUI(self):
        layout = QVBoxLayout()

        self.url_label = QLabel('M3U8 URL:')
        self.url_input = QLineEdit(self)
        layout.addWidget(self.url_label)
        layout.addWidget(self.url_input)

        self.cookies_label = QLabel('Cookies (可选):')
        self.cookies_input = QLineEdit(self)
        layout.addWidget(self.cookies_label)
        layout.addWidget(self.cookies_input)

        self.proxy_label = QLabel('代理 (可选):')
        self.proxy_input = QLineEdit(self)
        layout.addWidget(self.proxy_label)
        layout.addWidget(self.proxy_input)

        self.folder_label = QLabel('输出文件夹:')
        self.folder_input = QLineEdit(self)
        self.browse_button = QPushButton('浏览', self)
        self.browse_button.clicked.connect(self.browse_folder)
        layout.addWidget(self.folder_label)
        layout.addWidget(self.folder_input)
        layout.addWidget(self.browse_button)

        self.start_button = QPushButton('开始下载与合并', self)
        self.start_button.clicked.connect(self.on_start)
        layout.addWidget(self.start_button)

        self.progress_bar = QProgressBar(self)
        layout.addWidget(self.progress_bar)

        self.setLayout(layout)
        self.setWindowTitle('M3U8 下载器与合并工具')
        self.show()

    def browse_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "选择文件夹")
        if folder:
            self.folder_input.setText(folder)

    def on_start(self):
        self.start_button.setEnabled(False)
        m3u_url = self.url_input.text()
        proxy = self.proxy_input.text()
        output_folder = self.folder_input.text()
        cookies_text = self.cookies_input.text()
        cookies = {cookie.split('=')[0].strip(): cookie.split('=')[1].strip() for cookie in cookies_text.split(';') if '=' in cookie}

        self.thread = DownloadThread(m3u_url, output_folder, cookies, proxy)
        self.thread.progress_signal.connect(self.update_progress)
        self.thread.finished_signal.connect(self.enable_start_button)
        self.thread.start()

    def update_progress(self, current, total):
        self.progress_bar.setValue(int(((current+1) / (total+1)) * 100))

    def enable_start_button(self):
        self.start_button.setEnabled(True)

    def closeEvent(self, event):
        if self.thread is not None:
            if self.thread.isRunning:
                self.thread.request_stop()
                self.thread.wait()
        event.accept()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = App()
    sys.exit(app.exec_())
